import sys
import os
import argparse
import time
import numpy as np
import pynq
from multiprocessing import Process, Pipe, Queue, Event, Manager
from PIL import Image  # 如果PYNQ没有Pillow，可以用cv2
from IoU import *

parser = argparse.ArgumentParser(description="MobileNetv2 Tunable Activation Imbalance Transfer Quantization")
parser.add_argument('--f', type=int, default=333)
parser.add_argument('--v', type=int, default=850)
args = parser.parse_args()
print(args)

os.system("sudo sh -c \"sync && echo 3 > /proc/sys/vm/drop_caches\"")

BATCH_SIZE = 25
IMAGE_ROW, IMAGE_COL = 160, 320
N_THREAD = 8
IMG_DIR = '/home/huangjn/SkrSkr/sample1000/'  # 根据你的图片路径修改

def get_image_names():
    names_temp = [f for f in os.listdir(IMG_DIR) if f.endswith('.jpg')]
    names_temp.sort(key= lambda x:int(x[:-4]))
    return names_temp

def get_image_batch():
    image_list = get_image_names()
    batches = []
    for i in range(0, len(image_list), BATCH_SIZE):
        batches.append(image_list[i:i+BATCH_SIZE])
    return batches

def stitch(batch, image):
    for i in range(len(batch)):
        img_path = os.path.join(IMG_DIR, batch[i])
        img = Image.open(img_path).convert('RGB').resize((IMAGE_COL, IMAGE_ROW))
        image[i, :] = np.array(img).reshape(-1)
    return image

image_batches = get_image_batch()
N = len(image_batches) * BATCH_SIZE
assert N % BATCH_SIZE == 0

result = np.zeros(shape=(len(image_batches), BATCH_SIZE, 4), dtype=np.int32)
out_cnt = -1
ping, pong = 0, 1

img = [None] * 2
img[0] = pynq.allocate(shape=(BATCH_SIZE, IMAGE_ROW * IMAGE_COL * 3), dtype=np.uint8, cacheable=1)
img[1] = pynq.allocate(shape=(BATCH_SIZE, IMAGE_ROW * IMAGE_COL * 3), dtype=np.uint8, cacheable=1)
out = [None] * 2
out[0] = pynq.allocate(shape=(BATCH_SIZE, 2, 7), dtype=np.int16, cacheable=1)
out[1] = pynq.allocate(shape=(BATCH_SIZE, 2, 7), dtype=np.int16, cacheable=1)
print("Allocating memory done")

bitfile = "SkrSkr.bit"
overlay = pynq.Overlay(bitfile)

os.system("sudo dfs 0 "+str(args.f))
os.system("sudo dvs "+str(args.v))

dma = overlay.axi_dma
print("Bitstream loaded")



total_time = 0
print("******Start******")
power_list = []

for batch_idx, batch in enumerate(image_batches):
    # 加载图片到img[ping]
    stitch(batch, img[ping])
    start = time.time()
    dma.sendchannel.transfer(img[ping])
    dma.recvchannel.transfer(out[ping])
    
    dma.recvchannel.wait()
    dma.sendchannel.wait()
    end = time.time()
    total_time += end - start

    
    out_cnt += 1
    ping, pong = pong, ping


print("****Finished****")

length = (out_cnt + 1) * BATCH_SIZE
print("# of output: {:d}".format(length))
fps = length/total_time
print("Total time: {:f}s, FPS: {:f}".format(total_time, fps))

power = np.array(power_list)
avgPower = power.mean()
total_energy = avgPower * total_time
print("Total energy: {:f}mJ; Averge Power: {:f}mW".format(total_energy, avgPower))

result = result.reshape(-1, 4)
ground_truth_box = load_ground_truth('ground_truth.txt')
IoU = Average_IoU(ground_truth_box, result)

f = open("txt/{}.txt".format(args.v),"a+")
f.write("{}\t{}\t{:2f}\t{:2f}\t{:2f}\t{:2f}\n".format(args.v, args.f, fps, IoU*100, avgPower, total_energy))
f.close()
del img
print()